import os
import re
import sys
from collections import defaultdict
from concurrent.futures import ProcessPoolExecutor, as_completed
from dataclasses import dataclass, field
from pathlib import Path
from typing import Dict, List, Any, Tuple

import h5py
import numpy as np
from jsonargparse import ArgumentParser
from tqdm import tqdm

from mcap_base_reader import MCAPBaseReader, SyncedMCAPData


"""
使用例：
python data_asst/mcap_utils/mcap_2_hdf5.py \
    --task_name "task1 task2" \
    --input_dir "./input" \
    --output_dir "./output" \
    --max_workers 16 \
    --ignore_episodes "[[10, 13], [5, 8]]"
"""


@dataclass(frozen=True)
class ConversionConfig:
    """批量将 MCAP 文件转换为 HDF5 格式的配置"""

    # ... (这部分保持不变) ...
    task_name: str = field(
        default_factory=str
    )  # 需要处理的任务文件夹名称列表，用空格分隔
    ignore_episodes: List[List[int]] = field(
        default_factory=list
    )  # 每个任务忽略的episode列表
    input_dir: str = field(default_factory=str)  # 存储 MCAP 数据集的根输入路径
    output_dir: str = field(default_factory=str)  # 存储 HDF5 文件的根输出目录
    max_workers: int = os.cpu_count() // 2 or 1  # 并行处理的最大工作进程数

    def __post_init__(self):
        object.__setattr__(self, "input_dir", Path(self.input_dir))
        object.__setattr__(self, "output_dir", Path(self.output_dir))

        if isinstance(self.task_name, str):
            object.__setattr__(self, "task_name", self.task_name.split())

        if not self.ignore_episodes:
            object.__setattr__(self, "ignore_episodes", [[]] * len(self.task_name))


# 继承 MCAPBaseReader
class MCAPConverter(MCAPBaseReader):
    """
    处理单个 MCAP 文件，从基类获取同步数据，
    并将其保存为 HDF5 格式。
    """

    def __init__(self, mcap_file: Path, output_file: Path, task_name: str):
        # 初始化基类，传入 mcap_file
        super().__init__(mcap_file)
        self.output_file = output_file
        self.task_name = task_name
        # 注意：obs_fields, act_fields, camera_names 等已在基类中定义

    def convert(self) -> Dict[str, Any]:
        """执行转换流程，并返回结果字典"""
        result = {
            "input_file": str(self.mcap_file),
            "output_file": None,
            "success": False,
            "error": None,
        }
        try:
            # 1. 从基类获取同步后的数据
            synced_data = self.read_and_sync_data()

            if synced_data is None:
                result["error"] = "未找到有效的关节或图像数据（或同步失败）"
                print(
                    f"[!] 警告: 任务 '{self.task_name}' 的文件 {self.mcap_file.name} 因'{result['error']}'被跳过。"
                )
                return result

            # 2. 准备 HDF5 特定的数据集格式
            datasets, metadata = self._prepare_hdf5_datasets(synced_data)

            # 3. 保存到 HDF5
            self._save_to_hdf5(datasets, metadata)
            result.update({"output_file": str(self.output_file), "success": True})
            return result

        except Exception as e:
            result["error"] = str(e)
            print(
                f"[!] 错误: 转换任务 '{self.task_name}' 的文件 {self.mcap_file.name} 时发生异常: {e}"
            )
            return result

    def _prepare_hdf5_datasets(self, synced_data: SyncedMCAPData) -> Tuple[Dict, Dict]:
        """
        将 SyncedMCAPData 格式化为 HDF5 所需的字典。
        """

        # 基类已经返回了 [T, D] 格式的 numpy 数组
        qpos_list = synced_data.qpos_obs
        action_list = synced_data.qpos_act

        # 基类返回的 images 是 {"left_camera": [T,H,W,C], ...}
        # HDF5 需要的是 {"cam_left": [T,H,W,C], ...}

        images_dict_hdf5 = {}
        for hdf5_cam, mcap_cam_key in self.CAMERA_MAPPING.items():
            # 将帧列表 [T, H, W, C] 转换为一个 [T, H, W, C] 的 numpy 数组
            images_dict_hdf5[hdf5_cam] = np.array(synced_data.images[mcap_cam_key])

        datasets = {
            "/observations/qpos": qpos_list,
            "/action": action_list,
            **{
                f"/observations/images/{cam}": frames
                for cam, frames in images_dict_hdf5.items()
            },
        }
        metadata = {
            "num_timesteps": len(synced_data.timestamps),
            "camera_names": str(self.camera_names),
        }
        return datasets, metadata

    def _save_to_hdf5(self, datasets: Dict[str, np.ndarray], metadata: Dict):
        """(此方法保持不变)"""
        self.output_file.parent.mkdir(parents=True, exist_ok=True)
        with h5py.File(self.output_file, "w") as f:
            for name, data in datasets.items():
                is_image = "images" in name
                f.create_dataset(
                    name,
                    data=data,
                    compression="gzip",
                    compression_opts=4,
                    chunks=True,
                    dtype="uint8" if is_image else "float32",
                )
            for key, value in metadata.items():
                f.attrs[key] = value


# --- 工作函数和 main (保持不变) ---
def convert_file_worker(
    mcap_file: Path, output_file: Path, task_name: str
) -> Dict[str, Any]:
    """供多进程调用的工作函数, 处理单个mcap文件"""
    converter = MCAPConverter(mcap_file, output_file, task_name)
    return converter.convert()


def main():
    # ... (main 函数的其余部分保持不变) ...
    parser = ArgumentParser()
    parser.add_class_arguments(ConversionConfig, as_group=False)
    args = parser.parse_args()
    config = ConversionConfig(**vars(args))

    print(
        f"输入目录: {config.input_dir}\n输出目录: {config.output_dir}\n任务: {config.task_name}"
    )

    # 验证任务名和忽略列表的数量是否匹配
    if config.ignore_episodes and len(config.task_name) != len(config.ignore_episodes):
        print(
            f"[!] 错误: 提供了 {len(config.task_name)} 个任务名，但提供了 {len(config.ignore_episodes)} 个忽略列表。两者数量必须一致！"
        )
        sys.exit(1)

    jobs = []
    if not config.task_name:
        print("[!] 错误: 未指定任何 task_name。请使用 --task_name 提供任务名称。")
        return

    # 改造任务循环以实现过滤逻辑
    for task_idx, task in enumerate(config.task_name):
        task_input_dir = config.input_dir / Path(task)
        task_output_dir = config.output_dir / Path(task)

        # 获取当前任务对应的忽略列表 (使用集合以提高查找效率)
        ignored_indices = (
            set(config.ignore_episodes[task_idx]) if config.ignore_episodes else set()
        )
        if ignored_indices:
            print(
                f"[*] 信息: 任务 '{task}' 将忽略以下 episode 索引: {sorted(list(ignored_indices))}"
            )

        if not task_input_dir.is_dir():
            print(
                f"[*] 警告: 任务 '{task}' 的输入目录 {task_input_dir} 不存在，已跳过。"
            )
            continue

        mcap_files = list(task_input_dir.rglob("*.mcap"))

        def extract_episode_index(file_path: Path) -> int:
            match = re.search(r"(\d+)", file_path.stem)  # 或者用 r"episode_(\d+)"
            if match:
                return int(match.group(1))
            else:
                return float("inf")  # 无法匹配的排在最后

        mcap_files = sorted(mcap_files, key=extract_episode_index)

        for mcap_file in mcap_files:
            ep_index = extract_episode_index(mcap_file)
            if ep_index in ignored_indices:
                print(
                    f"[*] 跳过: 任务 '{task_input_dir.stem}' 的 episode_{ep_index} ({mcap_file.name}) 在忽略列表中。"
                )
                continue

            output_file = task_output_dir / f"episode_{ep_index}.hdf5"
            jobs.append((mcap_file, output_file, task_input_dir.stem))

    if not jobs:
        print(
            "[!] 错误: 在指定的任务目录中未找到任何需要转换的 .mcap 文件 (可能都被忽略了)。"
        )
        return

    print(
        f"共找到 {len(jobs)} 个mcap文件需要转换。使用 {config.max_workers} 个进程开始处理..."
    )

    all_results = []
    with ProcessPoolExecutor(max_workers=config.max_workers) as executor:
        futures = {
            executor.submit(
                convert_file_worker, mcap_path, hdf5_path, task_name
            ): mcap_path
            for mcap_path, hdf5_path, task_name in jobs
        }

        for future in tqdm(
            as_completed(futures), total=len(futures), desc="转换所有文件"
        ):
            try:
                all_results.append(future.result())
            except Exception as e:
                mcap_path = futures[future]
                print(f"[!] 错误: 处理文件 '{mcap_path.name}' 时发生严重错误: {e}")

    successful_conversions = [r for r in all_results if r["success"]]
    print("\n--- 转换完成 ---")
    print(
        f"总计处理: {len(all_results)} 个文件, 成功: {len(successful_conversions)} 个, 失败: {len(all_results) - len(successful_conversions)} 个"
    )


if __name__ == "__main__":
    main()
